package wtflog

import (
	"io"
)

// LevelWriter is a concept "copied" from logrus.
// We provide a builtin type wtfLevelWriter to implement this interface.
// Use NewLevelWriter in the logger if needed.
type LevelWriter interface {
	Write(p []byte) (int, error)
	WriteLevel(level Level, p []byte) (int, error)
}

// NewLevelWriter inits a level writer with defaultWriter and levelWriterMap.
// If levelWriterMap is nil or empty then all messages are logged to defaultWriter.
// If you only want to log all to stdout/stderr, use WithWriter(os.Stdout/os.Stderr) directly.
func NewLevelWriter(defaultWriter io.Writer, levelWriterMap map[Level]io.Writer) LevelWriter {
	return &wtfLevelWriter{defaultWriter, levelWriterMap}
}

type wtfLevelWriter struct {
	defaultWriter  io.Writer
	levelWriterMap map[Level]io.Writer
}

// Write implements io.Writer. Data will be written with lw.defaultWriter.
func (lw wtfLevelWriter) Write(p []byte) (int, error) {
	return lw.defaultWriter.Write(p)
}

// WriteLevel implements LevelWriter. Writer lw will first match argument level with level-writer map.
// If there's a matched writer, data will be written with that writer.
// Else use the default writer.
func (lw wtfLevelWriter) WriteLevel(level Level, p []byte) (int, error) {
	if lw.levelWriterMap != nil {
		if w, ok := lw.levelWriterMap[level]; ok {
			return w.Write(p)
		}
	}
	return lw.defaultWriter.Write(p)
}
